from modbus import *


class TerminalDevice(object):
    """
    定义了终端设备的通用方法
    """

    def __init__(self, sock, device_address):
        self.sock = sock
        self.address = device_address

    def _send(self, data: bytes):
        total_sent = 0
        while total_sent < len(data):
            sent = self.sock.send(data[total_sent:])
            if sent == 0:
                raise RuntimeError("socket connection broken")
            total_sent = total_sent + sent

    def _receive(self):
        return self.sock.recv(1024)


class Zs4gDtuDevice(TerminalDevice):
    """
    众山4g dtu设备
    """

    def is_alive(self):
        """
        判断设备是否在线
        收到b"\xaa\x55\x00\x04\xe0\x04\x00\xe8"后，
        设备会应答b"\xaa\x55\x00\x05\xe0\x04\x05\x00\xee"，
        由此判断是否在线
        """
        self._send(b"\xaa\x55\x00\x04\xe0\x04\x00\xe8")
        if self._receive() == b"\xaa\x55\x00\x05\xe0\x04\x05\x00\xee":
            return True
        else:
            return False


class ZsDtuRtuDevice(TerminalDevice):
    """
    众山dtu或rtu设备
    """

    #  DO寄存器列表
    do_index_list = [ZsModbusRegisterAddressEnum.DO1,
                     ZsModbusRegisterAddressEnum.DO2,
                     ZsModbusRegisterAddressEnum.DO3,
                     ZsModbusRegisterAddressEnum.DO4]

    # DI寄存器列表
    di_index_list = [ZsModbusRegisterAddressEnum.DI1,
                     ZsModbusRegisterAddressEnum.DI2,
                     ZsModbusRegisterAddressEnum.DI3,
                     ZsModbusRegisterAddressEnum.DI4]

    def set_do(self, index, status: ModbusWriteSingleCoilValueEnum):
        """
        设置do输出
        """
        if type(index) is int:
            index = ZsDtuRtuDevice.do_index_list[index]
        md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: self.address,
                                 ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.WRITE_SINGLE_COIL,
                                 ModbusStructureTypeEnum.REGISTER_START_ADDRESS: index,
                                 ModbusStructureTypeEnum.COIL_VALUE: status})
        self._send(md.get_blob())
        blob = self._receive()
        md = ModbusData(blob=blob)
        if md.get_decoded()[ModbusStructureTypeEnum.COMMAND] != ModbusCommandEnum.WRITE_SINGLE_COIL:
            raise RuntimeError("Response mismatch")

    def get_dos(self):
        """
        获取各个do的状态
        """
        md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: self.address,
                                 ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.READ_COIL,
                                 ModbusStructureTypeEnum.REGISTER_START_ADDRESS: ZsModbusRegisterAddressEnum.DO1,
                                 ModbusStructureTypeEnum.COIL_NUMBER: 2})
        self._send(md.get_blob())
        blob = self._receive()
        md = ModbusData(blob=blob)
        if md.decoded[ModbusStructureTypeEnum.COMMAND] != ModbusCommandEnum.READ_COIL:
            raise RuntimeError("Command mismatch")
        return md.decoded['coil0'], md.decoded['coil1'], md.decoded['coil2'], md.decoded['coil3']

    def get_dis(self):
        """
        获取各个di的状态
        """
        md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: self.address,
                                 ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.READ_DI,
                                 ModbusStructureTypeEnum.REGISTER_START_ADDRESS: ZsModbusRegisterAddressEnum.DI1,
                                 ModbusStructureTypeEnum.DI_NUMBER: 2})
        self._send(md.get_blob())
        blob = self._receive()
        md = ModbusData(blob=blob)
        if md.decoded[ModbusStructureTypeEnum.COMMAND] != ModbusCommandEnum.READ_DI:
            raise RuntimeError("Command mismatch")
        return md.decoded['di0'], md.decoded['di1'], md.decoded['di2'], md.decoded['di3']


class ZzRtuDevice(TerminalDevice):
    """
    继电器模块
    """
    do_index_list = [ZzModbusRegisterAddressEnum.DO1,
                     ZzModbusRegisterAddressEnum.DO2,
                     ZzModbusRegisterAddressEnum.DO3,
                     ZzModbusRegisterAddressEnum.DO4,
                     ZzModbusRegisterAddressEnum.DO5,
                     ZzModbusRegisterAddressEnum.DO6,
                     ZzModbusRegisterAddressEnum.DO7,
                     ZzModbusRegisterAddressEnum.DO8,
                     ZzModbusRegisterAddressEnum.DO9,
                     ZzModbusRegisterAddressEnum.DO10,
                     ZzModbusRegisterAddressEnum.DO11,
                     ZzModbusRegisterAddressEnum.DO12,
                     ZzModbusRegisterAddressEnum.DO13,
                     ZzModbusRegisterAddressEnum.DO14,
                     ZzModbusRegisterAddressEnum.DO15,
                     ZzModbusRegisterAddressEnum.DO16]

    def set_do(self, index, status: ModbusWriteSingleCoilValueEnum):
        """
        设置指定do的状态
        """
        if type(index) is int:
            index = ZzRtuDevice.do_index_list[index]
        md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: self.address,
                                 ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.WRITE_SINGLE_COIL,
                                 ModbusStructureTypeEnum.REGISTER_START_ADDRESS: index,
                                 ModbusStructureTypeEnum.COIL_VALUE: status})
        self._send(md.get_blob())
        blob = self._receive()
        md = ModbusData(blob=blob)
        if md.get_decoded()[ModbusStructureTypeEnum.COMMAND] != ModbusCommandEnum.WRITE_SINGLE_COIL:
            raise RuntimeError("Response mismatch")

    def get_dos(self):
        """
        获取各个do的状态
        """
        md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: self.address,
                                 ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.READ_COIL,
                                 ModbusStructureTypeEnum.REGISTER_START_ADDRESS: ZzModbusRegisterAddressEnum.DO1,
                                 ModbusStructureTypeEnum.COIL_NUMBER: 16})
        self._send(md.get_blob())
        blob = self._receive()
        md = ModbusData(blob=blob)
        if md.decoded[ModbusStructureTypeEnum.COMMAND] != ModbusCommandEnum.READ_COIL:
            raise RuntimeError("Command mismatch")
        return md.decoded['coil0'], md.decoded['coil1'], md.decoded['coil2'], md.decoded['coil3'], \
               md.decoded['coil4'], md.decoded['coil5'], md.decoded['coil6'], md.decoded['coil7'], \
               md.decoded['coil8'], md.decoded['coil9'], md.decoded['coil10'], md.decoded['coil11'], \
               md.decoded['coil12'], md.decoded['coil13'], md.decoded['coil14'], md.decoded['coil15']


class RaspDevice(TerminalDevice):
    """
    树莓派设备
    """

    #  REG寄存器列表
    reg_index_list = [RaspModbusRegisterAddressEnum.REG1]

    def get_regs(self):
        """
        获取REG1的状态
        """
        md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: self.address,
                                 ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.READ_INPUT_REGISTER,
                                 ModbusStructureTypeEnum.REGISTER_START_ADDRESS: RaspModbusRegisterAddressEnum.REG1,
                                 ModbusStructureTypeEnum.REGISTER_NUMBER: 1})
        print(md.get_blob())
        self._send(md.get_blob())
        blob = self._receive()
        md = ModbusData(blob=blob)
        if md.decoded[ModbusStructureTypeEnum.COMMAND] != ModbusCommandEnum.READ_INPUT_REGISTER:
            raise RuntimeError("Command mismatch")
        return md.decoded['reg1'], 0, 0, 0
